-- Copyright (c) 2007 Maksym Motornyy <mmotorny@gmail.com>
--
-- Permission is hereby granted, free of charge, to any person obtaining a copy
-- of this software and associated documentation files (the "Software"), to deal
-- in the Software without restriction, including without limitation the rights
-- to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-- copies of the Software, and to permit persons to whom the Software is
-- furnished to do so, subject to the following conditions:
--
-- The above copyright notice and this permission notice shall be included in
-- all copies or substantial portions of the Software.
--
-- THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-- IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-- FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-- AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-- LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-- OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
-- THE SOFTWARE.

module HPython.Lexer where

import Data.Char
import Data.Complex
import Text.ParserCombinators.Parsec hiding (notFollowedBy)
import HPython.ParsecExtensions

-- TODO: ignore identation inside all kinds of brackets

data Lexeme
	= NewLine
	| Indent
	| Dedent
	| Identifier String
	| Keyword String
	| BooleanLiteral Bool
	| StringLiteral String
	| IntegerLiteral Integer
	-- TODO:
	| FloatingPointLiteral Double
	-- TODO:
	| ImaginaryLiteral (Complex Double)
	| Operator String
	| Delimiter String
	| EndOfFile
	deriving Eq
	
instance Show Lexeme where
	show _lexeme =
		case _lexeme of
			NewLine -> "end of line"
			Indent -> "indent"
			Dedent -> "dedent"
			Identifier _name -> "identifier \"" ++ _name ++ "\""
			Keyword _name -> "keyword \"" ++ _name ++ "\""
			BooleanLiteral _value -> show _value
			IntegerLiteral _value -> show _value
			FloatingPointLiteral _value -> show _value
			ImaginaryLiteral _value -> show _value
			StringLiteral _value -> show _value
			Operator _name -> "operator \"" ++ _name ++ "\""
			Delimiter _name -> "delimiter \"" ++ _name ++ "\""
			EndOfFile -> "end of file"

data 
	LexemeAtPosition = 
	LexemeAtPosition
	{
		lexeme :: Lexeme,
		position :: SourcePos
	}

instance Show LexemeAtPosition where
	show _lexemeAtPosition =
		show $ lexeme _lexemeAtPosition

data 
	LexerState = 
	LexerState
	{
		indentationStack :: [Int],
		numberOfNestedBrackets :: Int
	}

type LexemeParser returnValue = GenParser LexemeAtPosition () returnValue

singleWhiteSpace :: CharParser LexerState Char
singleWhiteSpace =
	oneOf "\t\f "

whiteSpace :: CharParser LexerState [LexemeAtPosition]
whiteSpace =
	do
	{	many1 $ singleWhiteSpace
	;	return []
	}

comment :: CharParser LexerState [LexemeAtPosition]
comment =
	do
	{	char '#'
	;	many $ noneOf "\r\n"
	;	return []
	}
	
singleEndOfLine :: CharParser LexerState String
singleEndOfLine =
	do
	{	string "\r\n"
	<#>	string "\n"
	<#>	string "\r"
	}

lineJoining :: CharParser LexerState [LexemeAtPosition]
lineJoining =
	do
		char '\\'
		singleEndOfLine
		return []

indentation :: String -> Int
indentation _whiteSpaces =
	foldr
		(\_whiteSpace _indentation -> 
			case _whiteSpace of
				'\t' -> _indentation + 8
				_ -> _indentation + 1
		)
		0
		_whiteSpaces
		
endOfLine' :: SourcePos -> String -> SourcePos -> CharParser LexerState [LexemeAtPosition]
endOfLine' _newLinePosition _whiteSpaces _indentationPosition =
	do
	{	let _indentation = indentation _whiteSpaces
	;	_lexerState <- getState
	;	let _indentationStack = indentationStack _lexerState
	;	case compare _indentation (head _indentationStack) of
			LT ->
				if _indentation `elem` _indentationStack
					then do
						let _newIndentationStack =
							dropWhile (> _indentation) _indentationStack
						setState $
							_lexerState
							{
								indentationStack = _newIndentationStack
							}
						return $
							(LexemeAtPosition NewLine _newLinePosition) :
								(replicate 
									(length _indentationStack - length _newIndentationStack)
									(LexemeAtPosition Dedent _indentationPosition)
								)
					else
						pzero <?> "valid indentation at next line"
			EQ -> return [LexemeAtPosition NewLine _newLinePosition]
			GT -> do
				setState $
					_lexerState
					{
						indentationStack = _indentation:_indentationStack
					}
				return 
					[
						(LexemeAtPosition NewLine _newLinePosition), 
						(LexemeAtPosition Indent _indentationPosition)
					]
	}	
		
endOfLine :: CharParser LexerState [LexemeAtPosition]
endOfLine =
	do
		_newLinePosition <- getPosition
		singleEndOfLine
		_lexerState <- getState
		if numberOfNestedBrackets _lexerState > 0
			then return []
			else do
				_whiteSpaces <- many $ singleWhiteSpace
				_indentationPosition <- getPosition
				do
				{	do
						comment
						return [LexemeAtPosition NewLine _newLinePosition]
				<#>	do
						_endOfFile <- endOfFile
						return $ (LexemeAtPosition NewLine _newLinePosition):_endOfFile
				<#>	endOfLine
				<#>	endOfLine' _newLinePosition _whiteSpaces _indentationPosition
				}

identifier :: CharParser LexerState [LexemeAtPosition]
identifier =
	do
	{	_namePosition <- getPosition
	;	_name <-
			do
			{	firstCharacter <- (letter <|> char '_')
			;	restCharacters <- many $ (alphaNum <|> char '_')
			;	return $ firstCharacter:restCharacters
			}
	;	return [LexemeAtPosition (Identifier _name) _namePosition]
	}

keyword :: CharParser LexerState [LexemeAtPosition]
keyword =
	do
	{	_namePosition <- getPosition
	;	_name <-
			do
			{	string "and"
			<#>	string "assert"
			<#>	string "as"
			<#>	string "break"
			<#>	string "class"
			<#>	string "continue"
			<#>	string "def"
			<#>	string "del"
			<#>	string "elif"
			<#>	string "else"
			<#>	string "except"
			<#>	string "exec"
			<#>	string "finally"
			<#>	string "for"
			<#>	string "from"
			<#>	string "global"
			<#>	string "if"
			<#>	string "import"
			<#>	string "in"
			<#>	string "is"
			<#>	string "lambda"
			<#>	string "not"
			<#>	string "or"
			<#>	string "pass"
			<#>	string "print"
			<#>	string "raise"
			<#>	string "return"
			<#>	string "try"
			<#>	string "while"
			<#>	string "with"
			<#>	string "yield"
			}
	;	return [LexemeAtPosition (Keyword _name) _namePosition]
	}
	
booleanLiteral :: CharParser LexerState [LexemeAtPosition]
booleanLiteral =
	do
		_valuePosition <- getPosition
		_value <-
			do
			{	do
					string "False"
					return False
			<#>	do
					string "True"
					return True
			}
		return [LexemeAtPosition (BooleanLiteral _value) _valuePosition]

octalInteger :: CharParser LexerState Integer
octalInteger =
	do
	{	char '0'
	;	_digits <- many1 octDigit
	;	return $ read ("0o" ++ _digits)
	}

hexadecimalInteger :: CharParser LexerState Integer
hexadecimalInteger =
	do
	{	char '0'
	;	(oneOf "xX")
	;	_digits <- many1 hexDigit
	;	return $ read ("0x" ++ _digits)
	}

decimalInteger :: CharParser LexerState Integer
decimalInteger =
	do
	{	_digits <-
			do
			{	do
				{	firstDigit <- satisfy (`elem` ['1'..'9'])
				;	restDigits <- many digit
				;	return $ firstDigit:restDigits
				}
			<#>	string "0"
			}
	;	return $ read _digits
	}

integerLiteral :: CharParser LexerState [LexemeAtPosition]
integerLiteral =
	do
	{	_valuePosition <- getPosition
	;	_value <-
			do
			{	octalInteger
			<#>	hexadecimalInteger
			<#>	decimalInteger
			}
	;	option ' ' (oneOf "lL")
	;	return [LexemeAtPosition (IntegerLiteral _value) _valuePosition]
	}

integerPart :: CharParser LexerState Int
integerPart =
	do
		_digits <- many1 digit
		return $ read $ dropWhile (== '0') _digits

simpleFloatingPoint :: CharParser LexerState Double
simpleFloatingPoint =
	do
	{	do
			_integerPart <- option 0 integerPart
			char '.'
			_fractionalPart <- many1 digit
			return $
				fromIntegral _integerPart + (read $ "0." ++ _fractionalPart)
	<#>	do
			_integerPart <- integerPart
			char '.'
			return $ fromIntegral _integerPart
	}

exponentialFloatingPoint :: CharParser LexerState Double
exponentialFloatingPoint =
	do
		_base <-
			do
			{	simpleFloatingPoint
			<#>	do
					_integerPart <- integerPart
					return $ fromIntegral _integerPart
			}
		oneOf "eE"
		_operatorName <- option '+' (oneOf "+-")
		_signlessExponent <- integerPart
		let _exponent =
			case _operatorName of
				'+' -> _signlessExponent
				'-' -> -_signlessExponent
				_ -> error "Only '+' or '-' operator can be used in exponent"
		return $ _base ^^ _exponent

floatingPoint :: CharParser LexerState Double
floatingPoint =
	do
	{	exponentialFloatingPoint
	<#>	simpleFloatingPoint
	}

floatingPointLiteral :: CharParser LexerState [LexemeAtPosition]
floatingPointLiteral =
	do
		_valuePosition <- getPosition
		_value <- floatingPoint
		return [LexemeAtPosition (FloatingPointLiteral _value) _valuePosition]

imaginaryLiteral :: CharParser LexerState [LexemeAtPosition]
imaginaryLiteral =
	do
		_valuePosition <- getPosition
		_value <-
			do
			{	floatingPoint
			<#>	do
					_integerPart <- integerPart
					return $ fromIntegral _integerPart
			}
		oneOf "jJ"
		return [LexemeAtPosition (ImaginaryLiteral $ 0.0 :+ _value) _valuePosition]
	
-- TODO: \N{name} is not supported
unicodeEscapeSequence :: CharParser LexerState Char
unicodeEscapeSequence =
	do
	{	_prefix <- oneOf "uU"
	;	_code <- count (if _prefix == 'u' then 4 else 8) hexDigit
	;	return $ chr $ read ("0x" ++ _code)
	}

rawEscapeSequence :: CharParser LexerState Char -> CharParser LexerState String
rawEscapeSequence _unicodeEscapeSequence =
	do
	{	do
		{	singleEndOfLine
		;	return ['\\', '\n']
		}
	<#>	do
		{	_unicodeCharacter <- _unicodeEscapeSequence
		;	return [_unicodeCharacter]
		}
	<#>	do
		{	_character <- anyChar
		;	return ['\\', _character]
		}
	}

decodeEscapeSequence :: Char -> Char
decodeEscapeSequence _character =
	read ("'\\" ++ [_character] ++ "'")
	
normalEscapeSequence :: CharParser LexerState Char -> CharParser LexerState String
normalEscapeSequence _unicodeEscapeSequence =
	do
	{	do
		{	singleEndOfLine
		;	return ""
		}
	<#>	do
		{	_character <- oneOf "\\'\"abfnrtv"
		;	return [decodeEscapeSequence _character]
		}
	<#>	do
		{	_code <- count 3 octDigit
		;	return [chr $ read ("0o" ++ _code)]
		}
	<#>	do
		{	_code <- count 2 hexDigit
		;	return [chr $ read ("0x" ++ _code)]
		}
	<#>	do
		{	_unicodeCharacter <- _unicodeEscapeSequence
		;	return [_unicodeCharacter]
		}
	}

stringPrefix :: CharParser LexerState (CharParser LexerState String)
stringPrefix =
	do
	{	_prefix <-
			do
			{	string "ur"
			<#>	string "UR"
			<#>	string "Ur"
			<#>	string "uR"
			<#>	string "r"
			<#>	string "R"
			<#>	string "u"
			<#>	string "U"
			<#>	return ""
			}
	;	let _lowerCasePrefix = map toLower _prefix
	;	let _unicodeEscapeSequence =
			if 'u' `elem` _lowerCasePrefix
				then unicodeEscapeSequence
				else pzero
	;	if 'r' `elem` _lowerCasePrefix
			then return $ rawEscapeSequence _unicodeEscapeSequence
			else return $ normalEscapeSequence _unicodeEscapeSequence
	}

stringCharacters :: (CharParser LexerState String) -> (CharParser LexerState Char) -> String -> CharParser LexerState String
stringCharacters _escapeSequence _stringCharacter _delimiter =
	do
	{	notFollowedBy $ string _delimiter
	;	do
		{	do
			{	_character <- _stringCharacter
			;	return [_character]
			}
		<#>	do
			{	char '\\'
			;	_escapeSequence
			}
		}
	}
	
longStringCharacter :: CharParser LexerState Char
longStringCharacter =
	do
	{	do
		{	singleEndOfLine
		;	return '\n'
		}
	<#>	noneOf "\\"
	}
	
shortStringCharacter :: CharParser LexerState Char
shortStringCharacter =
	noneOf "\\\r\n"

stringLiteral' :: String -> (String -> CharParser LexerState String) -> CharParser LexerState String
stringLiteral' _delimiter _character =
	between 
		(string _delimiter) 
		(string _delimiter) 
		(do
			_characters <- many $ _character _delimiter
			return $ concat _characters
		)

stringLiteral :: CharParser LexerState [LexemeAtPosition]
stringLiteral =
	do
	{	_contentsPosition <- getPosition
	;	_escapeSequence <- stringPrefix
	;	let _stringCharacters = stringCharacters _escapeSequence
	;	_contents <-
			do
			{	stringLiteral' "\"\"\"" (_stringCharacters longStringCharacter)
			<#>	stringLiteral' "'''" (_stringCharacters longStringCharacter)
			<#>	stringLiteral' "\"" (_stringCharacters shortStringCharacter)
			<#>	stringLiteral' "'" (_stringCharacters shortStringCharacter)
			}
	; return [LexemeAtPosition (StringLiteral _contents) _contentsPosition]
	}

operator :: CharParser LexerState [LexemeAtPosition]
operator =
	do
	{	_namePosition <- getPosition
	;	_name <-
			do
			{	string "**"
			<#>	string "//"
			<#>	string "<<"
			<#>	string ">>"
			<#>	string "<="
			<#>	string ">="
			<#>	string "=="
			<#>	string "!="
			<#>	string "<>"
			<#>	string "+"
			<#>	string "-"
			<#>	string "*"
			<#>	string "/"
			<#>	string "%"
			<#>	string "&"
			<#>	string "|"
			<#>	string "^"
			<#>	string "~"
			<#>	string "<"
			<#>	string ">"
			}
	;	return [LexemeAtPosition (Operator _name) _namePosition]
	}

modifyNumberOfNestedBrackets :: (Int -> Int -> Int) -> CharParser LexerState ()
modifyNumberOfNestedBrackets _operator =
	do
		_lexerState <- getState
		setState $
			_lexerState
			{
				numberOfNestedBrackets = 
					numberOfNestedBrackets _lexerState `_operator` 1
			}
			
incrementNumberOfNestedBrackets :: CharParser LexerState ()
incrementNumberOfNestedBrackets =
	modifyNumberOfNestedBrackets (+)
	
decrementNumberOfNestedBrackets :: CharParser LexerState ()
decrementNumberOfNestedBrackets =
	modifyNumberOfNestedBrackets (-)
	
delimiter :: CharParser LexerState [LexemeAtPosition]
delimiter =
	do
	{	_namePosition <- getPosition
	;	_name <-
			do
			{	string "//="
			<#>	string ">>="
			<#>	string "<<="
			<#>	string "**="
			<#>	string "+="
			<#>	string "-="
			<#>	string "*="
			<#>	string "/="
			<#>	string "%="
			<#>	string "&="
			<#>	string "|="
			<#>	string "^="
			<#>	do
					_name <- string "("
					incrementNumberOfNestedBrackets
					return _name
			<#>	do
					_name <- string ")"
					decrementNumberOfNestedBrackets
					return _name
			<#>	do
					_name <- string "["
					incrementNumberOfNestedBrackets
					return _name
			<#>	do
					_name <- string "]"
					decrementNumberOfNestedBrackets
					return _name
			<#>	do
					_name <- string "{"
					incrementNumberOfNestedBrackets
					return _name
			<#>	do
					_name <- string "}"
					decrementNumberOfNestedBrackets
					return _name
			<#>	string "@"
			<#>	string ","
			<#>	string ":"
			<#>	string "."
			<#>	string "`"
			<#>	string "="
			<#>	string ";"
			}
	;	return [LexemeAtPosition (Delimiter _name) _namePosition]
	}
	
endOfFile :: CharParser LexerState [LexemeAtPosition]
endOfFile =
	do
	{	_endOfFilePosition <- getPosition
	;	eof
	;	_lexerState <- getState
	;	let _indentationStack = indentationStack _lexerState
	;	return $ 
				[LexemeAtPosition NewLine _endOfFilePosition]
			++	(replicate 
					(length _indentationStack - 1) 
					(LexemeAtPosition Dedent _endOfFilePosition)
				)
			++	[LexemeAtPosition EndOfFile _endOfFilePosition]
	}

removeDuplicateNewLines :: [LexemeAtPosition] -> [LexemeAtPosition]

removeDuplicateNewLines ((LexemeAtPosition NewLine _newLinePosition):_restLexemes) =
	(LexemeAtPosition NewLine _newLinePosition) :
	(removeDuplicateNewLines $ 
		dropWhile 
			(\_lexemeAtPosition -> lexeme _lexemeAtPosition == NewLine)
			_restLexemes
	)
		
removeDuplicateNewLines (_lexeme:_restLexemes) =
	_lexeme:(removeDuplicateNewLines _restLexemes)

removeDuplicateNewLines [] =
	[]

lexer :: CharParser LexerState [LexemeAtPosition]
lexer =
	do
		_lexemeLists <- 
			many $ 
				do
				{	whiteSpace
				<#>	comment
				<#> lineJoining
				<#>	endOfLine
				<#>	booleanLiteral
				<#>	imaginaryLiteral
				<#>	floatingPointLiteral
				<#>	integerLiteral
				<#>	stringLiteral
				<#>	operator
				<#>	delimiter
				<#>	keyword
				<#>	identifier
				}
		let _lexemes = concat _lexemeLists
		_endOfFile <- 
			if null _lexemes
				then endOfFile
				else
					case lexeme $ last _lexemes of
						EndOfFile -> return []
						_ -> endOfFile
		return $ removeDuplicateNewLines $ _lexemes ++ _endOfFile
